import time
import os
import cv2
import numpy as np
from PIL import Image
from deeplab import DeeplabV3

if __name__ == "__main__":

    new_model_path = "./best_epoch_weights.pth"  # 请将此路径替换为您的新模型文件路径
    deeplab = DeeplabV3(model_path=new_model_path)  # 在创建实例时传入新的模型路径

    
    mode = "predict"
    count = False
    name_classes = ["_background_", "leaf", "point"]
    
    if mode == "predict":
        save_path = "./predictions/"  # 保存预测结果的文件夹
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        while True:
            img = './vir322_1_img.jpg'  # 输入的图片路径
            try:
                image = Image.open(img)
            except Exception as e:
                print(f'Open Error! {e}. Try again!')
                continue
            else:
                r_image = deeplab.detect_image(image, count=count, name_classes=name_classes)
                
                # 保存预测结果，指定名称
                save_filename = os.path.join(save_path, os.path.basename(img))
                r_image.save(save_filename)
                print(f'Saved prediction to {save_filename}')
                
                # 显示结果图像
                # r_image.show()
            break  # 如果只想处理单张图片，处理完成后可退出循环，若需多张，移除此行

    # 其他模块 (如 video, fps, dir_predict 等) 保持不变
